//
// Created by 郝羽 on 2018/1/16.
//

#include <getopt.h>
#include <iostream>

extern "C" {
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/timestamp.h>
#include <libavutil/opt.h>
}
#include <map>
#include <opus.h>
#include <string>
#include <cmath>

#include "../../src/xtlv_file.h"
#include "../../src/xcutil.h"

enum Options {
    Help,
    FeedZero,
    Encode1,
    ExternalEncoder,
};


const char shortOptions[] = "h";

static const struct option longOptions[] = {
        {   "help",         no_argument,        nullptr,    Options::Help           },
        {   "feedzero",     no_argument,        nullptr,    Options::FeedZero       },
        {   "encode1",      no_argument,        nullptr,    Options::Encode1        },
        {   "ext",          no_argument,        nullptr,    Options::ExternalEncoder},
        {   0,              0,                  0,          0                       },
};

bool        feedZero = false;
bool        encodeOne = false;
bool        useExternalEncoder = false;
bool        disablePrediction = true;

void printUsage() {
    std::clog << "usage: tlv [options] input output" << std::endl;
    std::clog << "options:" << std::endl;
    std::clog << "       --help                  print this" << std::endl;
    std::clog << "       --feed-zero             feed zero to unused encoders, disabled by default" << std::endl;
}

//AVFormatContext* mixerOutputCtx = nullptr;
//std::map<int, AVCodecContext*> mixerOutputEncodeCtxs;
//AVFrame* mixerFrame = nullptr;
//AVFrame* mixerZeroFrame = nullptr;
//
//AVFormatContext* mixerInputCtx = nullptr;
//AVCodecContext* mixerInputEncodeCtx = nullptr;

AVFormatContext* newFile(const std::string& fileName, const std::string& formatName, const std::string codecName) {
    AVFormatContext* ctx = nullptr;
    int ret = avformat_alloc_output_context2(&ctx, nullptr, formatName.c_str(), fileName.c_str());
    if (ret != 0) {
        std::clog << "could not alloc output context : " << av_err2str(ret) << std::endl;
        return nullptr;
    }
    AVCodec* codec = avcodec_find_encoder_by_name(codecName.c_str());
    AVStream* stream = avformat_new_stream(ctx, codec);
    AVCodecParameters* par = stream->codecpar;

    par->format = AV_SAMPLE_FMT_S16;
    par->channels = 2;
    par->channel_layout = AV_CH_LAYOUT_STEREO;
    par->sample_rate = 16000;
    par->codec_id = codec->id;
    par->codec_type = codec->type;

    stream->time_base = AVRational {1, 16000};

    std::cout << "initialized output stream" << std::endl;
    av_dump_format(ctx, 0, nullptr, 1);

    ret = avio_open(&ctx->pb, fileName.c_str(), AVIO_FLAG_WRITE);
    if (ret < 0) {
        std::clog << "could not open file " << fileName << " : " << av_err2str(ret) << std::endl;
        avformat_free_context(ctx);
        return nullptr;
    }

    ret = avformat_write_header(ctx, nullptr);
    if (ret < 0) {
        std::clog << "could not write file header to " << fileName << " : " << av_err2str(ret) << std::endl;
        avformat_free_context(ctx);
        ctx = nullptr;
        return nullptr;
    }

    return ctx;
}
AVCodecContext* newEncoder(const std::string& codecName) {
    AVCodec* codec = avcodec_find_encoder_by_name(codecName.c_str());
    AVCodecContext* encCtx = avcodec_alloc_context3(codec);
    if (!encCtx) {
        std::clog << "could not alloc encoder ctx" << std::endl;
        return nullptr;
    }
    encCtx->codec_type = codec->type;
    encCtx->codec_id = codec->id;
    encCtx->sample_rate = 16000;
    encCtx->sample_fmt = AV_SAMPLE_FMT_S16;
    encCtx->channel_layout = AV_CH_LAYOUT_STEREO;
    encCtx->channels = 2;

    av_opt_set(encCtx->priv_data, "application", "voip", 0);

    int ret = avcodec_open2(encCtx, codec, nullptr);
    if (ret < 0) {
        std::clog << "could not open encode ctx : " << av_err2str(ret) << std::endl;
        avcodec_free_context(&encCtx);
        return nullptr;
    }
    return encCtx;
}

AVFrame* newFrame() {
    AVFrame* frame = av_frame_alloc();
    if (!frame) {
        std::clog << "could not alloc frame" << std::endl;
        return nullptr;
    }
    frame->format = AV_SAMPLE_FMT_S16;
    frame->channel_layout = AV_CH_LAYOUT_STEREO;
    frame->channels = 2;
    frame->sample_rate = 16000;
    frame->nb_samples = 20 * 16000 / 1000;
    int ret = av_frame_get_buffer(frame, 0);
    if (ret < 0) {
        std::clog << "could not alloc frame buffer" << std::endl;
        av_frame_free(&frame);
        return nullptr;
    }
    memset(frame->data[0], 0, frame->linesize[0]);

    return frame;
}

OpusEncoder* newOpusEncoder() {
    int err;
//    auto enc = opus_encoder_create(16000, 2, OPUS_APPLICATION_VOIP, &err);
    auto enc = opus_encoder_create(16000, 2, OPUS_APPLICATION_AUDIO, &err);
    if (err != 0) {
        std::clog << "create opus encoder error : " << std::string(opus_strerror(err)) << std::endl;
    }
    opus_encoder_ctl(enc, OPUS_SET_COMPLEXITY_REQUEST, 0);
    opus_encoder_ctl(enc, OPUS_SET_INBAND_FEC_REQUEST, 1);
    opus_encoder_ctl(enc, OPUS_SET_PACKET_LOSS_PERC_REQUEST, 20);

    if (disablePrediction) {
        int ret = opus_encoder_ctl(enc, OPUS_SET_PREDICTION_DISABLED(1));
        if (ret != 0) {
            std::clog << "could not disable prediction : " << std::string(opus_strerror(ret)) << std::endl;
        }
    }
    return enc;
}

void freeOpusEncoder(OpusEncoder* enc) {
    opus_encoder_destroy(enc);
}

int opusEncode(int64_t ts, AVFrame* frame, AVFormatContext* fmtCtx, OpusEncoder* enc) {

    uint8_t buf[1000];

    int ret = opus_encode(enc, (int16_t*)frame->data[0], frame->nb_samples, buf, sizeof(buf));
    if (ret < 0) {
        std::clog << "could not encode with external encoder : " << std::string(opus_strerror(ret)) << std::endl;
        return -1;
    }

    AVPacket packet;
    av_init_packet(&packet);
    packet.data = buf;
    packet.size = ret;
    packet.stream_index = 0;

    AVRational* time_base = &fmtCtx->streams[0]->time_base;
    std::clog << "pts:" << av_ts2str(packet.pts) << " ";
    std::clog << "pts-time:" << av_ts2timestr(packet.pts, time_base) << " ";
    std::clog << "duration:" << av_ts2str(packet.duration) << " ";
    std::clog << "duration-time:" << av_ts2timestr(packet.duration, time_base) << std::endl;
    ret = av_interleaved_write_frame(fmtCtx, &packet);
    if (ret < 0) {
        std::clog << "could not write frame : " << av_err2str(ret) << std::endl;
        return -2;
    }

    return 0;
}

int opusReinit(OpusEncoder* enc) {
    int ret = opus_encoder_init(enc, 16000, 2, OPUS_APPLICATION_VOIP);
    if (ret != 0) {
        std::clog << "could not re-init encoder : " << std::string(opus_strerror(ret)) << std::endl;
    }
    return ret;
}

int encode(int64_t& ts, AVFrame* frame, AVFormatContext* fmtCtx, AVCodecContext* encCtx) {
    int ret = avcodec_send_frame(encCtx, frame);
    if (ret < 0) {
        std::clog << "could not send frame to encoder : " << av_err2str(ret) << std::endl;
        return -2;
    }

    AVPacket packet;
    av_init_packet(&packet);
    while (avcodec_receive_packet(encCtx, &packet) == 0) {
        packet.pts = ts;
        packet.dts = packet.pts;
        ts += 20;
        AVRational* time_base = &fmtCtx->streams[0]->time_base;
        std::clog << "pts:" << av_ts2str(packet.pts) << " ";
        std::clog << "pts-time:" << av_ts2timestr(packet.pts, time_base) << " ";
        std::clog << "duration:" << av_ts2str(packet.duration) << " ";
        std::clog << "duration-time:" << av_ts2timestr(packet.duration, time_base) << std::endl;
        ret = av_interleaved_write_frame(fmtCtx, &packet);
        if (ret < 0) {
            std::clog << "could not write frame : " << av_err2str(ret) << std::endl;
            return -3;
        }
    }

    return 0;
}

int encodeAndThrowAway(AVFrame* frame, AVCodecContext* encCtx) {
    int ret = avcodec_send_frame(encCtx, frame);
    if (ret < 0) {
        std::clog << "could not feed zero frame to encoder : " << av_err2str(ret) << std::endl;
        return ret;
    }
    AVPacket packet;
    av_init_packet(&packet);
    while (avcodec_receive_packet(encCtx, &packet) == 0) {

    };
    return 0;
}

int opusEncodeAndThrowAway(AVFrame* frame, OpusEncoder* enc) {
    //opus_encoder_init(enc, 16000, 2, OPUS_APPLICATION_VOIP);
    uint8_t buf[1000];
    opus_encode(enc, (int16_t*)frame->data[0], frame->nb_samples, buf, sizeof(buf));
    return 0;
}

void freeFile(AVFormatContext* fmtCtx) {
    if (fmtCtx) {
        if (fmtCtx->pb) {
            av_write_trailer(fmtCtx);
            avio_close(fmtCtx->pb);
            fmtCtx->pb = nullptr;
        }
        avformat_free_context(fmtCtx);
    }
}

void freeEncoder(AVCodecContext* encCtx) {
    if (encCtx) {
        avcodec_free_context(&encCtx);
    }
}

void freeFrame(AVFrame* frame) {
    if (frame) {
        av_frame_free(&frame);
    }
}

int mixerLoop(const std::string& inputFileName) {
    FILE* in = fopen(inputFileName.c_str(), "rb");
    if (!in) {
        std::clog << "could not open file " << inputFileName << std::endl;
        return -1;
    }

    std::string inputFileNameBase(inputFileName);
    auto lastDot = inputFileName.rfind(".");
    inputFileNameBase = inputFileNameBase.substr(0, lastDot);

    std::string pcmOutFN = inputFileNameBase + "-pcm.wav";
    auto pcmOut = newFile(pcmOutFN, "wav", "pcm_s16le");
    if (!pcmOut) {
        std::clog << "could not new pcm output" << std::endl;
        return -1;
    }
    auto pcmEnc = newEncoder("pcm_s16le");
    if (!pcmEnc) {
        std::clog << "could not new pcm encoder" << std::endl;
        return -1;
    }
    int64_t pcmTs = 0;
    std::string opusOutFN = inputFileNameBase + "-opus.webm";
    auto opusOut = newFile(opusOutFN, "webm", "libopus");
    if (!opusOut) {
        std::clog << "could not new opus output" << std::endl;
        return -1;
    }
    auto opusEnc = newEncoder("libopus");
    if (!opusEnc) {
        std::clog << "could not new opus encoder" << std::endl;
        return -1;
    }
    int64_t opusTs = 0;

    OpusEncoder* extEnc = nullptr;
    if (useExternalEncoder) {
        extEnc = newOpusEncoder();
        if (!extEnc) {
            std::clog << "could not create external encoder" << std::endl;
            return -1;
        }
    }

    std::map<int, AVCodecContext*> opusEncMap;
    std::map<int, OpusEncoder*> extEncMap;

    auto frame = newFrame();
    auto zeroFrame = newFrame();

    uint8_t buf[4000];
    int len;
    int type;
    int last_type = -1;
    while (tlv_file_read_raw(in, buf, sizeof(buf), &len, &type) == 0) {
        bool encoder_changed = false;
        if (last_type != -1 && type != last_type) {
            encoder_changed = true;
            std::clog << "encoder switch " << last_type - 3000 << " -> " << type - 3000 << std::endl;
        }
        last_type = type;

        int ret = av_frame_make_writable(frame);
        if (ret < 0) {
            std::clog << "could not make frame writable : " << av_err2str(ret) << std::endl;
            return -1;
        }
        memcpy(frame->data[0], buf, len);

        if (encodeOne) {
            if (extEnc) {
                opusEncode(opusTs, frame, opusOut, extEnc);
            } else {
                encode(opusTs, frame, opusOut, opusEnc);
            }
        } else {
            if (extEnc) {
                auto it = extEncMap.find(type);
                if (it == extEncMap.end()) {
                    auto enc = newOpusEncoder();
                    if (!enc) {
                        std::clog << "could not create external opus encoder " << type << std::endl;
                        return -1;
                    }
                    auto insert = extEncMap.emplace(type, enc);
                    it = insert.first;
                }
                if (encoder_changed) {
//                    for (auto it : extEncMap) {
//                        if (it.first == type) {
//                            continue;
//                        }
//                    }
                    opusEncodeAndThrowAway(zeroFrame, it->second);
                }
                opusEncode(opusTs, frame, opusOut, it->second);
                memcpy(zeroFrame->data[0], frame->data[0], frame->linesize[0]);
//                if (feedZero) {
//                    for (auto it : extEncMap) {
//                        if (it.first == type) {
//                            continue;
//                        }
//                        opusEncodeAndThrowAway(zeroFrame, it.second);
//                    }
//                }
            } else {
                auto it = opusEncMap.find(type);
                if (it == opusEncMap.end()) {
                    auto enc = newEncoder("libopus");
                    if (!enc) {
                        std::clog << "could not new opus encoder " << type << std::endl;
                        return -1;
                    }
                    auto insert = opusEncMap.emplace(type, enc);
                    it = insert.first;
                }
                if (encoder_changed) {
                    encodeAndThrowAway(zeroFrame, it->second);
                }

                encode(opusTs, frame, opusOut, it->second);
                memcpy(zeroFrame->data[0], frame->data[0], frame->linesize[0]);
            }

        }

        double rms = 0.0;
        for (int i = 0; i < frame->nb_samples; i += 1) {
            int16_t* data = (int16_t*)(frame->data[0]);
            double level = (double)*data;
            level /= 32767;
            rms += std::pow(level, 2);
        }
        rms /= frame->nb_samples;
        rms = std::sqrt(rms);
        double level = 20.0 * std::log10(rms);
        if (level < -127.0) {
            level = -127.0;
        }
        //std::clog << "level = " << level << std::endl;
        for (int i = 0; i < frame->nb_samples; i += 1) {
            int16_t* data = (int16_t*)(frame->data[0]);
            //*(data + i * 2) = level;
            *(data + i * 2 + 1) = (type -3000) * 1000;
        }

        encode(pcmTs, frame, pcmOut, pcmEnc);

    }

    freeFrame(frame);

    freeOpusEncoder(extEnc);
    freeEncoder(opusEnc);
    for (auto it : opusEncMap) {
        freeEncoder(it.second);
    }
    freeFile(opusOut);

    freeEncoder(pcmEnc);
    freeFile(pcmOut);
}

int main(int argc, char* argv[]) {
    av_register_all();
    avcodec_register_all();

    int optionIndex = 0;
    int c;
    while((c = getopt_long(argc, argv, shortOptions, longOptions, &optionIndex)) != -1) {
        switch (c) {
            case Options::Help:
                printUsage();
                return 0;
            case Options::FeedZero:
                feedZero = true;
                break;
            case Options::Encode1:
                encodeOne = true;
                break;
            case Options::ExternalEncoder:
                useExternalEncoder = true;
                break;
            default:
                std::clog << "unknown option " << optopt << std::endl;
                return -1;
        }
    }

    if (argc < optind) {
        std::clog << "input file not specified" << std::endl;
        printUsage();
        return -2;
    }

    std::string inputFile = std::string(argv[optind]);
    std::cout << "setting input file " << inputFile << std::endl;
    optind += 1;

    if (argc > optind) {
        for (int i = optind; i < argc; ++i) {
            std::clog << "unknown option found : " << std::string(argv[i]) << std::endl;
        }
        printUsage();
        return -3;
    }

    mixerLoop(inputFile);
}